"""
@Author: Bright
@File: doc_loader_custom.py
@Time: 2025/9/30
@Desc:
"""
import asyncio
from typing import Iterator

from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document
from onnx.gen_proto import qualify


class CustomDocLoader(BaseLoader):
    """从文件逐行读取的示例文档加载器。 """

    def __init__(self, file_path: str):
        """
        使用文件路径初始化加载起。
        :param file_path:要加载的文件的路径。
        """
        self.file_path = file_path

    def lazy_load(self) -> Iterator[Document]:
        """逐行读取文件的惰性加载器。
        当你实现惰性加载方法时，应使用生成器，逐个生成文档。
        :return:
        """
        with open(self.file_path, "r", encoding="utf-8") as f:
            line_number = 0
            for line in f:
                # yield关键字用于定义生成器函数，
                # 生成器函数是一种特殊类型的函数，它允许你逐步生成一个序列的值，而不是一次返回整个序列。
                # 与普通的函数不同，生成器在每次调用时，不会从头开始执行，而是从上次离开的地方继续执行。
                # 这使得生成器非常适合处理需要[逐步生成或处理大数据集]的情况。
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1

    # alazy_load是可选的。如果你省略了实现，将使用默认实现，改实现将委托给lazy_load
    async def async_load(self) -> Iterator[Document]:
        """逐行读取文件的【异步】惰性加载器。"""
        # 需要aiofiles, 使用pip install aiofiles 安装。 https://github.comTinche/aiofiles
        import aiofiles
        async with aiofiles.open(self.file_path, encoding="utf-8") as f:
            line_number = 0
            async for line in f:
                # async for 逐行异步读取文件内容，并通过yield关键字逐行生成Document对象。
                # 异步生成器允许你在异步环境中逐步生成值，这对处理I/O密集型任务（如文件读取和网络请求）非常有用。
                # 因为它可以在等待I/O操作完成时，释放时间循环，以便处理其他任务。
                yield Document(
                    page_content=line,
                    metadata={"line_number": line_number, "source": self.file_path},
                )
                line_number += 1


with open("./mew.txt", "w", encoding="utf-8") as f:
    quality_content = "喵喵🐱 \n 喵喵🐱 \n 喵喵🐱"
    f.write(quality_content)
loader = CustomDocLoader("./mew.txt")

## 测试延迟加载接口
for doc in loader.lazy_load():
    print()
    print(type(doc))
    print(doc)


# 测试异步实现
async def alazy():
    async for doc in loader.async_load():
        print(type(doc))
        print(doc)


print("--" * 50)
asyncio.run(alazy())

"""
在生产环境中，应该避免使用load()，一次行加载所有数据到内存中,影响性能.
"""
print(loader.load())
